import matplotlib.pyplot as plt
from collections import Counter
import json
import jieba.posseg as pseg
# # 假设我们有以下文本数据
# text_data = [
# “他正在跑步”,
# “她喜欢唱歌”,
# “我们决定去旅行”,
# "我正在学习编程",
# “他们计划开公司”
# ]


def get_text_data(jsonl_file_path):
    text_data = []
    with open(jsonl_file_path, 'r', encoding='utf-8') as file:
        for line in file:
            try:
                data = json.loads(line)
                line_str = data['instruction']
                text_data.append(line_str)
            except json.JSONDecodeError:
                print(f"Error decoding JSON from line: {line}")
    return text_data


# 函数：提取句子中的动词
def extract_verbs(sentence):
    verbs = []
    # 对句子进行分词和词性标注
    words = pseg.cut(sentence)
    for word, flag in words:
        # 根据词性标注，提取动词（这里是动词和动名词）
        if flag == 'v' or flag == 'vn':
            verbs.append(word)
    return verbs


# 函数：提取文本中的主要动词
def extract_main_verbs(sentences):
    verbs = []
    for sentence in sentences:
        sentence_verbs = extract_verbs(sentence)
        verbs.extend(sentence_verbs)
    return verbs


# 函数：生成饼状图
def generate_pie_chart(data):
    verbs_count = Counter(data)
    labels = verbs_count.keys()
    sizes = verbs_count.values()
    colors = plt.cm.tab20c.colors  # 选择颜色映射

    # 设置忽略的最小百分比阈值
    min_percent_threshold = 1

    # 计算每个部分的百分比
    total_size = sum(sizes)
    percents = [size / total_size * 100 for size in sizes]

    # 移除低于阈值的部分
    visible_sizes = [size for size, percent in zip(sizes, percents) if percent >= min_percent_threshold]
    visible_labels = [label for label, percent in zip(labels, percents) if percent >= min_percent_threshold]

    plt.pie(visible_sizes, labels=visible_labels, colors=colors, autopct='%1.1f%%', startangle=140)
    plt.axis('equal')  # 确保饼图是圆形的
    plt.title('主要动词分布')


if __name__ == '__main__':
    # 设置matplotlib支持汉字
    plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用宋体
    plt.rcParams['axes.unicode_minus'] = False  # 正确显示负号
    jsonl_file_path = '../mydata/kimi_generations/qa_pairs_with_score.jsonl'
    text_data = get_text_data(jsonl_file_path)
    main_verbs = extract_main_verbs(text_data)
    generate_pie_chart(main_verbs)
    plt.show()